# -*- coding:utf-8 -*-
from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import FileResponse
import psutil
import speedtest
import os
import shutil
import uvicorn
import zipfile
import time
from datetime import datetime
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

'''
启动
pip install fastapi psutil speedtest-cli "uvicorn[standard]" python-multipart
uvicorn main:app --reload  --host 0.0.0.0 --port 7079
'''


'''
虚拟机网速测试
'''
@app.get("/test-internet-speed")
async def test_internet_speed():
    try:
        test = speedtest.Speedtest()
        test.get_servers()
        test.get_best_server()
        download_speed = float(test.download() / 1024)
        upload_speed = float(test.upload() / 1024)
    except Exception as e:
        #print("error s" % e)
        return {
            "code": 201,
            "data": {}
        }
    else:
        return {
            "code": 200,
            "data": {
                'download_speed': round(download_speed, 2),
                'upload_speed': round(upload_speed, 2),
            }
        }

'''
虚拟机磁盘空间详情
'''
@app.get("/df-h")
async def dh_f():
    try:
        partitions = psutil.disk_partitions()
        df = []
        for partition in partitions:
            if not os.path.exists(partition.mountpoint):
                continue

            usage = psutil.disk_usage(partition.mountpoint)
            df.append({
                'device': partition.device,
                'total': round(usage.total / (1024 ** 3), 2),
                'used': round(usage.used / (1024 ** 3), 2),
                'free': round(usage.free / (1024 ** 3), 2)
            })
    except Exception as e:
        #print("error s" % e)
        return {
            "code": 201,
            "data": {},
        }
    else:
        return {
            "code": 200,
            "data": {
                'df_h': df,
                'total': {
                    'device': "&".join(list(map(lambda x: x['device'], df))),
                    'total': round(sum(list(map(lambda x: x['total'], df))), 2),
                    'used': round(sum(list(map(lambda x: x['used'], df))), 2),
                    'free': round(sum(list(map(lambda x: x['free'], df))), 2)
                }
            }
        }

'''
虚拟机桌面文件中转站中的文件列表
'''
@app.get("/filetransfer-station-files")
async def get_file_transfer_station_files():
    _, file_transfer_station_files = _get_desktop_path()
    return {
        "code": 200,
        "data": {
            'file_transfer_station_files': file_transfer_station_files
        }
    }


'''
从虚拟机桌面文件中转站中下载文件
'''
@app.get("/file-download")
async def file_download(fileName: str):
    files = fileName.split(",")
    filePath, zip_name = _zip_files(files)
    return FileResponse(
        filePath,
        **{"filename": f"{zip_name}"}
    )

'''
上传文件到虚拟机桌面文件中转站
'''
@app.post("/file-upload")
async def file_upload(file: UploadFile =  File(...)):
    file_transfer_station_path, _ = _get_desktop_path()
    file_path = os.path.join(file_transfer_station_path, file.filename)
    # 判断是否存在同名文件
    if os.path.exists(file_path):
        # 生成当前时间的字符串，以当前时间为基础
        time_string = datetime.now().strftime("%Y%m%d%H%M%S")
        # 在文件名后面添加当前时间的字符串
        file_name, file_ext = os.path.splitext(file.filename)
        new_file_name = f"{file_name}_{time_string}{file_ext}"
        file_path = os.path.join(file_transfer_station_path, new_file_name)
    with open(file_path, "wb") as f:
        shutil.copyfileobj(file.file, f)
    return {
        "code": 200,
        "data": {
            "fileName": file.filename,
            "filePath": file_path,
            "contentType": file.content_type
        }
    }

def _get_desktop_path(stationName='FileTransferStation', file_formats=['.txt', '.doc', '.docx', '.xls', '.pdf']):   # in ubuntu is FileTransferStation
    desktop_path = os.path.join(os.path.expanduser("~"), 'Desktop')
    file_transfer_station_path = os.path.join(desktop_path, stationName)
    file_transfer_station_files = [
        file for file in os.listdir(file_transfer_station_path)
        if os.path.isfile(os.path.join(file_transfer_station_path, file))
        and any(file.endswith(file_format) for file_format in file_formats)
    ]
    return file_transfer_station_path, file_transfer_station_files

# def _get_desktop_path(stationName = '文件中转站'):             # in ubuntu is FileTransferStation
#     desktop_path = os.path.join(os.path.expanduser("~"), 'Desktop')
#     file_transfer_station_path = os.path.join(desktop_path, stationName)
#     file_transfer_station_files = [file for file in os.listdir(file_transfer_station_path) if os.path.isfile(os.path.join(file_transfer_station_path, file))]
#     return file_transfer_station_path, file_transfer_station_files

# def _zip_files(files):
#     file_transfer_station_path, _ = _get_desktop_path()
#     total_original_size = 0  # 初始化原文件大小总和为0
#     if len(files) == 1:
#         file_path = os.path.join(file_transfer_station_path, files[0])
#         return file_path, files[0], os.path.getsize(file_path) / (1024 * 1024)  # 返回文件路径、文件名和大小（MB）
#     files_to_compress = [os.path.join(file_transfer_station_path, file) for file in files]
#     zip_name = f"{time.time()}.zip"
#     zip_path = os.path.join(file_transfer_station_path, zip_name)
#     with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zip_ref:
#         for file_name in files_to_compress:
#             original_file_size = os.path.getsize(file_name)
#             total_original_size += original_file_size  # 累加原文件大小
#             zip_ref.write(file_name, arcname=file_name.replace(file_transfer_station_path, "files"))
#     total_original_size_mb = total_original_size / (1024 * 1024)  # 将原文件大小之和转换为MB
#     return zip_path, zip_name, total_original_size_mb

def _zip_files(files):
    file_transfer_station_path, _ = _get_desktop_path()
    if len(files) == 1:
        return os.path.join(file_transfer_station_path, files[0]), files[0]
    files_to_compress = [os.path.join(file_transfer_station_path, file) for file in files]
    zip_name = f"{time.time()}.zip"
    zip_path = os.path.join(file_transfer_station_path, zip_name)
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zip_ref:
        for file_name in files_to_compress:
            zip_ref.write(file_name, **{"arcname": file_name.replace(file_transfer_station_path, "files" )})
    return zip_path, zip_name


# in ubuntu use, in windows use start.bat
if __name__ == '__main__':
    uvicorn.run('main:app', host='0.0.0.0', port=7079, log_level="info")


